# @Time : 2021/8/6 10:04
# @Author : Li Kunlun
# @Description : 语言模型数据集(歌词)

from mxnet import nd
import random
import zipfile

# 1、读取数据集
with zipfile.ZipFile('../data/jaychou_lyrics.txt.zip') as zin:
    with zin.open('jaychou_lyrics.txt') as f:
        corpus_chars = f.read().decode('utf-8')
"""
查看前40个字符：
    想要有直升机
    想要和你飞到宇宙去
    想要和你融化在一起
    融化在宇宙里
    我每天每天每
"""
print(corpus_chars[:40])

# 为了打印方便，将换行符换成空格
corpus_chars = corpus_chars.replace('\n', ' ').replace('\r', ' ')
# 仅用前一万个字符训练模型
corpus_chars = corpus_chars[0:10000]

# 2、建立字符索引
idx_to_char = list(set(corpus_chars))
# <class 'dict'>: {'蓝': 0, '慢': 1,
char_to_idx = dict([(char, i) for i, char in enumerate(idx_to_char)])
vocab_size = len(char_to_idx)
# 1027
print(vocab_size)

# 将训练级中每个字符转化为索引
corpus_indices = [char_to_idx[char] for char in corpus_chars]
sample = corpus_indices[:20]
# chars: 想要有直升机 想要和你飞到宇宙去 想要和
# indices: [896, 512, 139, 550, 702, 815, 183, 896, 512, 320, 805, 164, 36, 179, 293, 438, 183, 896, 512, 320]
print('chars:', ''.join([idx_to_char[idx] for idx in sample]))
print('indices:', sample)


# 3、时序数据的采样
# 3.1、随机采样
def data_iter_random(corpus_indices, batch_size, num_steps, ctx=None):
    """
    :param
        :param corpus_indices: 每个字符索引
        :param batch_size: 批量大小
        :param num_steps: 时间步数
        :param ctx:
        :return:
    """
    # 减1是因为输出的索引是相应输入的索引加1
    num_examples = (len(corpus_indices) - 1) // num_steps
    epoch_size = num_examples // batch_size
    example_indices = list(range(num_examples))
    random.shuffle(example_indices)

    # 返回从pos开始的长为num_steps的序列
    def _data(pos):
        return corpus_indices[pos: pos + num_steps]

    for i in range(epoch_size):
        # 每次读取batch_size个随机样本
        i = i * batch_size
        batch_indices = example_indices[i: i + batch_size]
        X = [_data(j * num_steps) for j in batch_indices]
        Y = [_data(j * num_steps + 1) for j in batch_indices]
        yield nd.array(X, ctx), nd.array(Y, ctx)


my_seq = list(range(30))
for X, Y in data_iter_random(my_seq, batch_size=3, num_steps=6):
    """
    X:  
    [[ 18.  19.  20.  21.  22.  23.]
     [  0.   1.   2.   3.   4.   5.]]
    <NDArray 2x6 @cpu(0)> 
    Y: 
    [[ 19.  20.  21.  22.  23.  24.]
     [  1.   2.   3.   4.   5.   6.]]
    <NDArray 2x6 @cpu(0)> 
    
    X:  
    [[ 12.  13.  14.  15.  16.  17.]
     [  6.   7.   8.   9.  10.  11.]]
    <NDArray 2x6 @cpu(0)> 
    Y: 
    [[ 13.  14.  15.  16.  17.  18.]
     [  7.   8.   9.  10.  11.  12.]]
    <NDArray 2x6 @cpu(0)> 
    """
    print('X: ', X, '\nY:', Y, '\n')


# 3.1、相邻采样
def data_iter_consecutive(corpus_indices, batch_size, num_steps, ctx=None):
    corpus_indices = nd.array(corpus_indices, ctx=ctx)
    data_len = len(corpus_indices)
    batch_len = data_len // batch_size
    indices = corpus_indices[0: batch_size * batch_len].reshape((batch_size, batch_len))
    epoch_size = (batch_len - 1) // num_steps
    for i in range(epoch_size):
        i = i * num_steps
        X = indices[:, i: i + num_steps]
        Y = indices[:, i + 1: i + num_steps + 1]
        yield X, Y


for X, Y in data_iter_consecutive(my_seq, batch_size=2, num_steps=6):
    """
    X:  
    [[  0.   1.   2.   3.   4.   5.]
     [ 15.  16.  17.  18.  19.  20.]]
    <NDArray 2x6 @cpu(0)> 
    Y: 
    [[  1.   2.   3.   4.   5.   6.]
     [ 16.  17.  18.  19.  20.  21.]]
    <NDArray 2x6 @cpu(0)> 
    
    X:  
    [[  6.   7.   8.   9.  10.  11.]
     [ 21.  22.  23.  24.  25.  26.]]
    <NDArray 2x6 @cpu(0)> 
    Y: 
    [[  7.   8.   9.  10.  11.  12.]
     [ 22.  23.  24.  25.  26.  27.]]
    <NDArray 2x6 @cpu(0)> 
    """
    print('X: ', X, '\nY:', Y, '\n')
